return_source = ["news_title_zh", "news_content_zh", "news_summary", "news_publicdate"]
if __name__ == "__main__":
    import json

    from elasticsearch import Elasticsearch

    from args import *

    # 连接到Elasticsearch集群
    client = Elasticsearch(hosts=ELASTIC_ADDRESS, basic_auth=(ELASTIC_USERNAME, ELASTIC_PASSWORD), timeout=100)
    from model import *

    model = SentenceTransformer(MODEL_NAME, cache_folder=MODEL_PATH)
    from trans import *

    first_vector = model.encode(final)
    print(final)
    query = {

        # "query": {
        #     "match": {
        #         "news_content_zh": {
        #             "query": " ".join(keywords[:5]),
        #             "boost": 1
        #         }
        #     }},
        "knn": {
            "field": vector_name,
            "query_vector": first_vector,
            "k": 10,
            "num_candidates": 20,
            "boost": 1

        },
        "min_score": 0.94,
        "_source": return_source,
        "size": 10,

        # ,
        # "sort": [
        #     {
        #         "news_publicdate": {
        #             "order": "desc"
        #         }
        #     }
        # ],
    }

    # 输出query
    # print("查询向量的query：")
    # print(query.__str__().replace("'", '"'))
    # 输出查询结果
    import time

    start = time.time()
    result = client.search(index=INDEX_NAME, body=query, timeout="100s")
    print("查询向量结果：")
    result = result.body['hits']['hits']

    result = [i for i in result if i["_score"] < 0.99]
    score_buck = {}
    for r in result:
        cut = f"{r['_score']:.6f}"
        if cut not in r:
            score_buck[cut] = [r]
        else:
            score_buck[cut].append(r)
    result = [score_buck[k][0] for k in score_buck]

    print(f"查询时间：{time.time() - start: .2f}s")
    print(json.dumps(result, ensure_ascii=False, indent=4))
    # 保存结果到result.json
    with open("result.json", "w", encoding="utf-8") as f:
        # print(json.dumps([i["_source"] for i in result], ensure_ascii=False, indent=4))
        # json.dump([i["_source"] for i in result], f, ensure_ascii=False, indent=4)

        json.dump(result, f, ensure_ascii=False, indent=4)

    print("查询向量的结果已被保存到result.json")
